use crate::{
    RegistryError,
    models::{Entity, EntityRef},
};
use std::{
    collections::{BTreeMap, HashMap, HashSet},
    sync::Arc,
};
use uuid::Uuid;

const SERVICE_NAME_INDEX: &'static str = "name";
const SERVICE_GROUP_INDEX: &'static str = "group";

pub struct Index {
    inner: HashMap<&'static str, BTreeMap<String, HashSet<EntityRef>>>,
}

impl Index {
    pub fn new() -> Self {
        // create default index keys
        let inner = [
            (SERVICE_GROUP_INDEX, BTreeMap::new()),
            (SERVICE_NAME_INDEX, BTreeMap::new()),
        ]
        .into_iter()
        .collect();

        Index { inner }
    }

    pub fn insert_into(&mut self, service: EntityRef) {
        let index_table = &mut self.inner;

        // insert into specific name tree
        index_table
            .get_mut(SERVICE_NAME_INDEX)
            .unwrap()
            .entry(service.name.clone())
            // if key existed, insert value into vec
            .and_modify(|v| {
                let _ = v.insert(EntityRef(Arc::clone(&service)));
            })
            // if key not existe, insert key with single vec that wrapped with service
            .or_insert_with(|| {
                let mut hashset = HashSet::new();
                let _ = hashset.insert(EntityRef(Arc::clone(&service)));
                hashset
            });

        // insert into specific group tree
        index_table
            .get_mut(SERVICE_GROUP_INDEX)
            .unwrap()
            .entry(service.group.clone())
            // if key existed, insert value into vec
            .and_modify(|v| {
                let _ = v.insert(EntityRef(Arc::clone(&service)));
            })
            // if key not existe, insert key with single vec that wrapped with service
            .or_insert_with(|| {
                let mut hashset = HashSet::new();
                let _ = hashset.insert(EntityRef(Arc::clone(&service)));
                hashset
            });
    }

    pub fn remove_index(&mut self, entity: &EntityRef) {
        self.inner
            // query specific name tree from hashmap
            .get_mut(SERVICE_NAME_INDEX)
            .unwrap()
            // query reference of service entiry from btree
            .get_mut(&entity.name)
            .unwrap()
            .remove(&entity);

        self.inner
            // query specific name tree from hashmap
            .get_mut(SERVICE_GROUP_INDEX)
            .unwrap()
            // query reference of service entiry from btree
            .get_mut(&entity.group)
            .unwrap()
            .remove(&entity);
    }

    pub fn filter_by_name<'k>(&mut self, key: &'k str) -> Option<&HashSet<EntityRef>> {
        self.inner
            // query specific name tree from hashmap
            .get(SERVICE_NAME_INDEX)
            .unwrap()
            // query reference of service entiry from btree
            .get(key)
    }

    pub fn filter_by_group<'k>(&mut self, key: &'k str) -> Option<&HashSet<EntityRef>> {
        self.inner
            // query specific name tree from hashmap
            .get(SERVICE_GROUP_INDEX)
            .unwrap()
            // query reference of service entiry from btree
            .get(key)
    }
}

pub struct State {
    entity: BTreeMap<uuid::Uuid, EntityRef>,
    index: Index,
}

impl State {
    pub fn new() -> Self {
        State {
            entity: BTreeMap::new(),
            index: Index::new(),
        }
    }

    pub fn all_entities(&self) -> Option<Vec<EntityRef>> {
        let entity_set: Vec<EntityRef> =
            self.entity.values().map(|e| EntityRef::clone(e)).collect();
        if entity_set.is_empty() {
            None
        } else {
            Some(entity_set)
        }
    }

    pub fn register(&mut self, service: Entity) -> Result<(), RegistryError> {
        let uid = service.uid;
        // determine whether the entity already existes, if so return err, or continue
        if self.entity.contains_key(&uid) {
            return Err(RegistryError::AlreadyExist);
        }
        // service.
        let service_ref = Arc::new(service);
        // insert into btree first
        self.entity.insert(uid, EntityRef::clone(&service_ref));
        // insert into index
        self.index.insert_into(EntityRef::clone(&service_ref));

        Ok(())
    }

    pub fn unregister(&mut self, uid: Uuid) -> Result<(), RegistryError> {
        // remove the entity first
        self.entity
            .remove(&uid)
            .and_then(|entity| {
                // remove from index
                self.index.remove_index(&entity);
                Some(())
            })
            .ok_or(RegistryError::Empty)
    }

    #[inline]
    pub fn query_with_name<'a>(&mut self, name: &'a str) -> Option<&HashSet<EntityRef>> {
        self.index.filter_by_name(name)
    }

    #[inline]
    pub fn query_with_group<'a>(&mut self, group: &'a str) -> Option<&HashSet<EntityRef>> {
        self.index.filter_by_group(group)
    }

    #[inline]
    pub fn query_with_uid(&mut self, uid: Uuid) -> Option<EntityRef> {
        self.entity
            .get(&uid)
            .and_then(|e| Some(EntityRef::clone(&e)))
    }
}

#[cfg(test)]
mod test {
    use std::net::Ipv4Addr;

    use uuid::Uuid;

    use crate::db::{Entity, State};

    #[test]
    fn test_db() {
        let services = [
            Entity {
                uid: Uuid::new_v4(),
                host: std::net::IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
                port: 8081,
                name: "client-server-01".into(),
                group: "client".into(),
                tag: None,
            },
            Entity {
                uid: Uuid::new_v4(),
                host: std::net::IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
                port: 80,
                name: "db-server-01".into(),
                group: "client".into(),
                tag: None,
            },
            Entity {
                uid: Uuid::new_v4(),
                host: std::net::IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
                port: 80,
                name: "db-server-01".into(),
                group: "client".into(),
                tag: None,
            },
            Entity {
                uid: Uuid::new_v4(),
                host: std::net::IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
                port: 80,
                name: "user-service".into(),
                group: "db".into(),
                tag: None,
            },
        ];

        let mut db = State::new();

        services.into_iter().for_each(|s| {
            let _ = db.register(s);
        });

        // let same_names = db.query_with_name("user-service");
        let same_group = db.query_with_group("client");

        println!("{:?}", same_group);
    }
}
